# third party
import os
import sys
import Queue
import struct
import urllib
import string
import socket
import random
import threading

# first party
import bencoding
from misc import sha1
from misc import threadPrint
from peer import Peer
from fragmentedPiece import FragmentedPiece

class Torrent:

	def __init__ (self, owner, torrentFileName = None):

		# misc initialization
		self.owner = owner
		self.lock = threading.Lock()
		self.peers = []

		# this is only used in the end-game
		self.requestMasterList = []

		# this is kept up-to-date, but I haven't started using it at all yet
		self.piecePopulation = []

		# if we can init right away from a file, we do it!
		if torrentFileName != None:
			self.initFromFile( torrentFileName )

		


	def initFromFile (self, torrentFileName):
		self.alive = True
		self.endgame = False
	
		# populate metaInfo with the contents of the torrent file
		torrentFile = open( torrentFileName, "rb" ).read()
		torrentChars = list( torrentFile )
		torrentChars.reverse()
		self.metaInfo = bencoding.decodeItem( torrentChars )

		# unpack the piece hashes from the metaInfo
		self.pieceHash = []
		for i in range(0, len(self.metaInfo["info"]["pieces"]), 20):
			self.pieceHash.append( self.metaInfo["info"]["pieces"][i:i+20] )
		
		self.pieceCount = len(self.pieceHash)
		
		# the sha1 hash of metaInfo["info"] is used as a GUID for the torrent
		self.GUID = sha1( bencoding.bencodeDict(self.metaInfo["info"]) )

		# the handshake is used by peers to open communications
		protocolString = "BitTorrent protocol"
		self.handshake = struct.pack( "!B", len(protocolString) )
		self.handshake += struct.pack( str(len(protocolString)) + "s" , protocolString )
		self.handshake += struct.pack( "8x" )
		self.handshake += struct.pack( "20s", self.GUID )
		self.handshake += struct.pack( "20s", self.owner.localData["id"] )

		# let's deal with the target files!
		self.files = []
		self.fileSizes = []

		# if we're a multi-file torrent, read all the files in like so
		if "files" in self.metaInfo["info"]:
			for item in self.metaInfo["info"]["files"]:
				folderList = [self.metaInfo["info"]["name"]] + item["path"]
				fileSize = item["length"]
				self.createDummyFile( folderList, fileSize )
				self.files.append( string.join( folderList, "/" ) )
				self.fileSizes.append( fileSize )

		# single-file torrents stash the data a little differently, is all
		else:
			folderList = [self.metaInfo["info"]["name"]]
			fileSize = self.metaInfo["info"]["length"]
			self.createDummyFile( folderList, fileSize )
			self.files.append( string.join( folderList, "/" ) )
			self.fileSizes.append( fileSize )
			

		# miscelaneous initialization
		self.pieceComplete = []
		for i in range(self.pieceCount):
			self.pieceComplete.append( False )
		self.refreshBitfield()
		
		self.uploaded = 0
		self.downloaded = 0
		self.pieceFragments = {}

		# this is whildly inneficient, but it'll work
		# it's three in the morning, and I don't have a lot of time left.
		self.scanPieces()

		# peers need to be refreshed, can't be taken from saved data
		for i in range(self.pieceCount):
			if self.pieceComplete[i]:
				self.piecePopulation.append(sys.maxint)
			else:
				self.piecePopulation.append(0)
		

	# I'm pretty sure that I'm doing this wrong 
	# 	- should be able to just pickle the whole torrent in one go
	#	- will switch to this later, with research and experimentation
	#	- right now, I need to get some fucking data flowing already
	def getSaveData (self):
		return {"GUID" : self.GUID,
				"handshake" : self.handshake,
				"metaInfo" : self.metaInfo,
				"files" : self.files,
				"fileSizes" : self.fileSizes,
				"pieceCount" : self.pieceCount,
				"pieceHash" : self.pieceHash,
				"pieceComplete" : self.pieceComplete,
				"bitfield" : self.bitfield,
				"uploaded" : self.uploaded,
				"downloaded" : self.downloaded,
				"remaining" : self.remaining, 
				"pieceFragments" : self.pieceFragments,
				"missingPieces" : self.missingPieces }


	def initFromSavedData (self, savedData):
		self.GUID = savedData["GUID"]
		self.handshake = savedData["handshake"]
		self.metaInfo = savedData["metaInfo"]
		self.files = savedData["files"]
		self.fileSizes = savedData["fileSizes"]
		self.pieceCount = savedData["pieceCount"]
		self.pieceHash = savedData["pieceHash"]
		self.pieceComplete = savedData["pieceComplete"]
		self.bitfield = savedData["bitfield"]
		self.uploaded = savedData["uploaded"]
		self.downloaded = savedData["downloaded"]
		self.remaining = savedData["remaining"]

		self.pieceFragments = savedData["pieceFragments"]
		self.missingPieces = savedData["missingPieces"]

		self.alive = True
		self.endgame = False

		# this needs to be refreshed each time- can't reload it.
		for i in range(self.pieceCount):
			self.piecePopulation.append(sys.maxint if self.pieceComplete[i] else 0)

		for piece in self.pieceFragments:
			self.pieceFragments[piece].refreshFragmentsNeeded()

		random.shuffle(self.missingPieces)


	def getPieceSize (self, piece):
		if piece == self.pieceCount - 1:
			return sum(self.fileSizes) - ((self.pieceCount - 1) * self.metaInfo["info"]["piece length"])
		return self.metaInfo["info"]["piece length"]

	def getFragmentedPiece (self, piece):
		if not piece in self.pieceFragments:
			self.pieceFragments[piece] = FragmentedPiece( self.getPieceSize(piece) )
		return self.pieceFragments[piece]

		
	def contactTracker (self, event = None):
		url = self.metaInfo["announce"]
		
		values = { 	"info_hash" : self.GUID,
					"peer_id" : self.owner.localData["id"],
					"port" : self.owner.localPort,
					"uploaded" : self.uploaded,
					"downloaded" : self.downloaded,
					"left" : self.remaining,
					"compact" : 1,
					"no_peer_id" : 0 }

		if event:
			values["event"] = event
		
		data = urllib.urlencode(values)

		bencodedResponse = list( urllib.urlopen( url + "?" + data ).read() )
		bencodedResponse.reverse()
		
		response = bencoding.decodeItem( bencodedResponse )
		self.readPeers( response['peers'] )


	def readPeers (self, packedPeers):
		for i in range(0, len(packedPeers), 6):
			# unpack the peer data
			ipTuple = struct.unpack("!BBBB", packedPeers[i:i+4])
			peerIP = string.join( map(str, ipTuple), "." )
			peerPort = struct.unpack("!H", packedPeers[i+4:i+6])[0]

			Peer(self, ip=peerIP, port=peerPort)
			#break
			


	def receiveFragment (self, piece, offset, fragment):
		with self.lock:
			if self.pieceComplete[piece]:
				return
		
			fragmentedPiece = self.getFragmentedPiece(piece)
			fragmentedPiece.addFragment( offset, fragment )

			if self.endgame:
				request = [piece, offset, len(fragment)]
				if request in self.requestMasterList:
					self.requestMasterList.remove( request )
				for peer in self.peers:
					peer.cancelRequest( request )

			if fragmentedPiece.isComplete():
				pieceData = fragmentedPiece.getData()
				
				if sha1(pieceData) == self.pieceHash[piece]:
					self.missingPieces.remove( piece )
					self.downloaded += len(pieceData)
					self.pieceComplete[piece] = True
					self.piecePopulation[piece] = sys.maxint
					self.refreshBitfield()
					self.writePiece( piece, pieceData )
					del self.pieceFragments[piece]

					for peer in self.peers:
						peer.sendHave(piece)

				# if the hash didn't match, we screwed up the download somehow
				else:
					del self.pieceFragments[piece]
					
			


	def getFragment (self, piece, offset, length):
		with self.lock:
			if self.pieceComplete[piece]:
				return self.readPiece(piece)[offset:offset + length]
			return None

	

	def createDummyFile (self, pathList, fileSize):
		filePath = ""

		# lay down the folders, if they don't exist yet
		for i in range(len(pathList) - 1):
			if len(filePath) > 0:
				filePath += "/"
			filePath += pathList[i]
			
			if not os.path.isdir( filePath ):
				os.mkdir( filePath )

		# lay down the file, if it doesn't exist yet
		if len(filePath) > 0:
			filePath += "/"
		filePath += pathList[len(pathList)-1]

		if os.path.isfile(filePath):
			dummyfile = open( filePath, "ab" )
		else:
			dummyfile = open( filePath, "wb" )

		# resize the file to spec
		if os.path.getsize(filePath) != fileSize:
			dummyfile.truncate(fileSize - 1)
			dummyfile.seek(fileSize - 1)
			dummyfile.write("x")
		dummyfile.close()


	def scanPieces (self):
		self.missingPieces = []
		self.remaining = sum(self.fileSizes)
		
		for piece in range(self.pieceCount):
			pieceData = self.readPiece(piece)
			self.pieceComplete[piece] = ((sha1(pieceData) == self.pieceHash[piece]))

			if self.pieceComplete[piece]:
				self.remaining -= self.getPieceSize(piece)
			else:
				self.missingPieces.append(piece)

		random.shuffle(self.missingPieces)
		self.refreshBitfield()


	def readPiece (self, piece):
		offset = piece * self.metaInfo["info"]["piece length"]
		size = self.getPieceSize(piece)

		# which file do we start in?
		fileIndex = 0;
		while self.fileSizes[fileIndex] < offset:
			offset -= self.fileSizes[fileIndex]
			fileIndex += 1

		# read data in from the file
		data = ""
		while len(data) < size:
			chunkfile = open(self.files[fileIndex], "rb")
			chunkfile.seek(offset)
			data += chunkfile.read( size - len(data) )
			chunkfile.close()

			# if we aren't done, we go to the front of the next file
			fileIndex += 1
			offset = 0
			
		return data


	def writePiece (self, piece, data):
		offset = piece * self.metaInfo["info"]["piece length"]

		fileIndex = 0;
		while self.fileSizes[fileIndex] < offset:
			offset -= self.fileSizes[fileIndex]
			fileIndex += 1

		while len(data) > 0:
			writeSize = min( self.fileSizes[fileIndex] - offset, len(data) )
		
			chunkfile = open(self.files[fileIndex], "r+b")
			chunkfile.seek(offset)
			chunkfile.write( data[:writeSize] )
			chunkfile.close
			
			data = data[writeSize:]

			# if we aren't done, we go to the front of the next file
			fileIndex += 1
			offset = 0

			

	def refreshBitfield (self):
		self.bitfield = ""
		for i in range(0, self.pieceCount, 8):
			byte = 0
			for j in range(8):
				if i + j >= self.pieceCount:
					break;
				if self.pieceComplete[i + j]:
					byte = byte | (128 >> j)
			self.bitfield += struct.pack( "B", byte )
		self.owner.saveData()



	def nextFragmentRequest (self, peer):
		with self.lock:
			if not self.checkEndgame():
			
				# first, we try to get a request from any of the pieces we're already working on
				for piece in self.pieceFragments:
					if peer.pieceComplete[piece]:
						request = self.pieceFragments[piece].nextFragmentRequest()
						if request != None:
							offset, size = request
							return [piece, offset, size]

				# if we didn't find anything that we can use in the shortlist, look in the missingPieces list
				# note that I move through the missing pieces from back to front- this is an efficiency thing
				for i in range(len(self.missingPieces)):
					piece = self.missingPieces[ len(self.missingPieces) - i - 1 ]
				
					if peer.pieceComplete[piece]:
						request = self.getFragmentedPiece(piece).nextFragmentRequest()
						if request != None:
							offset, size = request
							return [piece, offset, size]

			return None
		

	def requestFailed (self, request):
		if self.endgame and self.alive and request in self.requestMasterList:
			for peer in self.peers:
				peer.sendRequest( request )
		else:
			piece, offset, size = request
			self.getFragmentedPiece(piece).requestFailed([offset, size])


	def checkEndgame (self):
		if self.endgame:
			return True
	
		# are we down to the last ten missing pieces?
		if len(self.missingPieces) > 10:
			return False
			
		# do any of the missing pieces have unsent requests?
		for missingPiece in self.missingPieces:
			fragmentedPiece = self.getFragmentedPiece(missingPiece)
			if fragmentedPiece.countUnsentRequests() > 0:
				return False

		# if we made it this far, it's party time
		self.startEndgame()
		return True
			
			

	def startEndgame (self):

		#threadPrint("---------------------------------------------------------------------------------------", 
		#			"\n-----------------------------------END GAME, BITCHES-----------------------------------",
		#			"\n---------------------------------------------------------------------------------------" )
	
		self.endgame = True
		self.requestMasterList = []

		for piece in self.missingPieces:
			fragmentedPiece = self.getFragmentedPiece(piece)
			
			for request in fragmentedPiece.getAllFragmentRequests():
				offset, size = request		
				self.requestMasterList.append( [piece, offset, size] )

		for peer in self.peers:
			random.shuffle( self.requestMasterList )
			for request in self.requestMasterList:
				peer.sendRequest( request )
		

	def reportUpload (self, upload):
		with self.lock:
			self.uploaded += upload

	def incrementPiecePopulation (self, piece):
		with self.lock:
			if not self.pieceComplete[piece]:
				self.piecePopulation[piece] += 1

	def decrementPiecePopulation (self, piece):
		with self.lock:
			if not self.pieceComplete[piece]:
				self.piecePopulation[piece] = min( 0, self.piecePopulation[piece] - 1)

	def killPeer (self, peer):
		with self.lock:	
			if peer in self.peers:
				self.peers.remove( peer )
				for i in range(self.pieceCount):
					if peer.pieceComplete[i] and not self.pieceComplete[i]:
						self.piecePopulation[i] = min( 0, self.piecePopulation[i] - 1)
				for request in peer.activeRequests:
					self.requestFailed( request )
				peer.queuedRequests = []

			

	def addPeer (self, peer):
		with self.lock:
			self.peers.append( peer )

			if self.endgame:
				random.shuffle( self.requestMasterList )
				for request in self.requestMasterList:
					peer.sendRequest( request )
				peer.updateActiveRequests()
			

	def kill (self):
		self.alive = False
		for peer in self.peers:
			peer.kill()



		
